#!/bin/env python3
#swapoff!
import ithildin as ith
import numpy as np
import matplotlib.pyplot as plt
import sys
from typing import List
from scipy.interpolate import RectBivariateSpline
#im
from pydiffmap.diffusion_map import DiffusionMap
from pydiffmap import visualization as diff_visualization
import matplotlib

data = ith.SimData.from_stem("../KORRR/phase_1000/phase_1000")
dmap = DiffusionMap.from_sklearn(n_evecs = 1, alpha=0.5, epsilon = 'bgh', k=4)

# de u en v coordinaaten voor een zekere (z,y,x)-positie
uc = data.vars['u'][:,:,200,200].flatten()
uv = data.vars['v'][:,:,200,200].flatten()

matplotlib.use("TkAgg")
plt.scatter(uc,uv)
plt.show()

faseruimte = np.array([uc,uv]) # bekend figuur! lijkt 1-dimensionaal (behalve voor rechtsboven eventueel, maar toch bij benadering ten minste)

# O jee geen convergentie
# dmap.fit(faseruimte.transpose())

# misschien kan convergentie opgelost worden met meer data?

uc = data.vars['u'].flatten()
uv = data.vars['v'].flatten()

# er zijn te veel punten voor plt.scatter om binnen redelijke tijd
# te werken, dus neem een willekeurige selectie

keuzes = np.random.choice(np.arange(uc.shape[0]), 60000,replace=False)
uc2 = uc[keuzes]
uv2 = uv[keuzes]
matplotlib.use("TkAgg")
plt.scatter(uc2, uv2)
plt.show() # dit lijkt er meer op!  Opnieuw kan er een soort eendimensionale structuur ontdekt worden maar toch wel wat afwijking.

# Kan dmap.fit hiermee iets doen? (verlaag aantal punten voor snelheid)
keuzes = np.random.choice(np.arange(uc.shape[0]), 1000,replace=False)
uc2 = uc[keuzes]
uv2 = uv[keuzes]
faseruimte2 = np.array([uc2, uv2])
dmap.fit(faseruimte2.transpose())
# (beëindigd wegens te lang bezig)

# misschien is het probleem dat de faseruimte niet echt in een tweedimensionale
# Euclidische ruimte past (en de invoer was al tweedimensionaal).  Voorstel:
# hoogdimensionaal werken.

# Returns: an array of shape (N, X×Y×Z)
# where (X,Y,Z) are the number of X, Y and Z positions respectively.
def variable(data, var):
    v = data.vars[var]
    return np.reshape(v,(v.shape[0], v.shape[1]*v.shape[2]*v.shape[3]))

def all_variables(data,vars):
    # Shape: (number of time steps, number of vars * number of grid elements)
    return np.append(*map(lambda var: variable(data,var),vars),axis=1)

phase_space = all_variables(data,['u','v'])

dmap = DiffusionMap.from_sklearn(n_evecs=2)
dmap.fit(phase_space) # gelukt!

matplotlib.use("TkAgg")
diff_visualization.data_plot(dmap,dim=3)
diff_visualization.embedding_plot(dmap,dim=2)


# uittesten meer dimensies

dmap = DiffusionMap.from_sklearn(n_evecs=3)
dmap.fit(phase_space) # gelukt!

matplotlib.use("TkAgg")
diff_visualization.data_plot(dmap,dim=3)
diff_visualization.embedding_plot(dmap,dim=3)



data_BOCF = ith.SimData.from_stem("../KORRR/results_9010/PDL_9010")

# Eerst even kijken
var1 = data_BOCF.vars['u'].ravel() # ravel: no copy
var2 = data_BOCF.vars['v'].ravel()
var3 = data_BOCF.vars['w'].ravel()
var4 = data_BOCF.vars['s'].ravel() #???? waarom faalt de code onder nu?
keuzes = np.random.choice(var1.shape[0], 5000,replace=True) # replace=True strikt genomen incorrect maar verlaagt geheugengebruik en wegens de hoge hoeveelheid punten niet belangrijk
var1_ = var1[keuzes]
var2_ = var2[keuzes]
var3_ = var3[keuzes]
var4_ = var4[keuzes]
matplotlib.use("TkAgg")
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(var1_, var2_, var3_, c=var4_,cmap=plt.hot())
plt.show()
del keuzes

#for k = seq(1,var1.size

# oeps twee verre groepjes punten (misschien interfereren die later?)
eigenaardig = (var3_ > 0.9999999) | (var1_ > 0.5)
#eigenaardig = (var3_ > 0.9999999) | (var1_ > 0.5) # ander groepje
var1__ = np.ma.masked_where(eigenaardig, np.ma.array(var1_))
var2__ = np.ma.masked_where(eigenaardig, np.ma.array(var2_))
var3__ = np.ma.masked_where(eigenaardig, np.ma.array(var3_))
var4__ = np.ma.masked_where(eigenaardig, np.ma.array(var4_))

fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(var1__, var2__, var3__,c=var4__,cmap=plt.hot())
plt.show()


del var1, var2, var3, var4, var1_, var2_, var3_, var4_, var1__, var2__, var3__, var4__, fig, ax
# oeps




def weight_negeer_uitschieters(datum):
#    s = data.vars['w'].shape
    print(datum)
    return oeps()

dmap_BOCF = DiffusionMap.from_sklearn(n_evecs=4,k=32, weight_fxn=weight_negeer_uitschieters) # k kleiner voor snelheid


# data_BOCF is een beetje te groot om mee te werken dus maken we
# het een stuk kleiner zodat np.append op mijn computer ermee overweg kan.

def verkleinde_variable(data, var, step_23):
    v = data.vars[var]
    points2 = np.arange(0,v.shape[2],step_23)
    points3 = np.arange(0,v.shape[3],step_23)
    vverkleind = v[:,:,points2,:]
    vverkleind = v[:,:,:,points3]
    reshaped = np.reshape(vverkleind,(vverkleind.shape[0], vverkleind.shape[1]*vverkleind.shape[2]*vverkleind.shape[3]))
    return reshaped

def verkleinde_all_variables(data,vars):
    # Shape: (number of time steps, number of vars * reduced number of grid elements)
    # TODO: grootte skip
    #   ok: 50, 38, 31
    #  nok: 40
    return np.concatenate(list(map(lambda var: verkleinde_variable(data,var,31),vars)),axis=1)

dmap_BOCF.fit(verkleinde_all_variables(data_BOCF, ['u', 'v', 'w', 's']))

matplotlib.use("TkAgg")
#diff_visualization.data_plot(dmap_BOCF,dim=3)

gc.collect()
# Plot volgens de tweede, derde en vierde coordinaat. Kleur volgens de eerste.
matplotlib.use("TkAgg")
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(dmap_BOCF.dmap[:,1],dmap_BOCF.dmap[:,2],dmap_BOCF.dmap[:,3],c=dmap_BOCF.dmap[:,0], cmap='viridis')
plt.show()


# kleur volgens de vierde nieuwe coördinaat
scatter_kwarg = {'c' : dmap_BOCF.dmap[:,3] }
diff_visualization.embedding_plot(dmap_BOCF,dim=3,scatter_kwargs=scatter_kwarg)
#diff_visualization.data_plot(dmap_BOCF,dim=3,n_evecs=1)


gc.collect()
# Plot volgens ???
matplotlib.use("TkAgg")
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
ax.scatter(dmap_BOCF.dmap[:,3],dmap_BOCF.dmap[:,1],dmap_BOCF.dmap[:,2],c=dmap_BOCF.dmap[:,0], cmap='viridis')
plt.show()



# ! globaal
# ! hoe werkt pydiffmap, ph1,...2 in verband brengen met algoritme,
#   theorie uitleggen van plotjes (halve pagina puntjes)

# <!> plotje BOFC (buento-...) PDL_9010



## Laten we eens opnieuw proberen, deze keer lokaal

step_23_ = 50
def verkleinde_variable2(data, var):
    v = data.vars[var]
    points2 = np.arange(0,v.shape[2],step_23_)
    points3 = np.arange(0,v.shape[3],step_23_)
    vverkleind = v[:,:,points2,:]
    vverkleind = v[:,:,:,points3]
    return vverkleind

data_BOCF = ith.SimData.from_stem("../KORRR/results_9010/PDL_9010")
def weight_negeer_uitschieters(datum):
    s = data.vars['w'].shape
    print(s)
    print(datum)
    return oeps()
dmap_BOCF = DiffusionMap.from_sklearn(n_evecs=4,k=32) #, weight_fxn=weight_negeer_uitschieters)
dmap_BOCF.fit(np.array((verkleinde_variable2(data_BOCF, 'u').flatten(),verkleinde_variable2(data_BOCF, 'v').flatten(),verkleinde_variable2(data_BOCF, 'w').flatten(),verkleinde_variable2(data_BOCF, 's').flatten())).transpose())

#dmap_BOCF.fit(verkleinde_all_variables(data_BOCF, ['u', 'v', 'w', 's']))

